import pandas as pd
import torch
import yaml
from torch_geometric.data import Data
import numpy as np
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
from torch import nn
from torch_geometric.datasets import Planetoid
import random


def set_seed(seed=0):
    random.seed(seed)  # Python
    np.random.seed(seed)  # NumPy
    torch.manual_seed(seed)  # CPU
    torch.cuda.manual_seed(seed)  # GPU
    torch.cuda.manual_seed_all(seed)  # GPU
    torch.backends.cudnn.deterministic = True  #
    torch.backends.cudnn.benchmark = False  #


def read_yaml(file_path: str) -> dict:
    # 打开并读取 YAML 文件
    with open(file_path, 'r', encoding='utf-8') as file:
        content = file.read()
    data = yaml.safe_load(content)
    return data


def construct_graph_data(df: pd.DataFrame) -> Data:
    """
    根据传入的pandas构建图数据
    :param df:
    :return:
    """
    data = Data()
    data.x = torch.from_numpy(np.vstack(df.features.values)).float()
    edge_index = []
    for user_id in df.id.values:
        for page_id in df.loc[user_id, "action"]:
            edge_index.append((user_id, page_id))
    data.edge_index = torch.tensor(edge_index).t()
    data.y = torch.tensor(df.label.values, dtype=torch.int64)
    data.train_mask = torch.zeros(len(df), dtype=torch.bool)
    data.valid_mask = torch.zeros(len(df), dtype=torch.bool)
    data.test_mask = torch.zeros(len(df), dtype=torch.bool)
    data.train_mask[torch.from_numpy(df.loc[df.flag == "Train"].index.values)] = True
    data.valid_mask[torch.from_numpy(df.loc[df.flag == "Valid"].index.values)] = True
    data.test_mask[torch.from_numpy(df.loc[df.flag == "Test"].index.values)] = True
    return data


def construct_graph_data_hetero(df: pd.DataFrame) -> dict:
    """
    根据传入的pandas构建图数据
    :param df:
    :return:
    """
    edge_index = []
    for user_id in df.id.values:
        for page_id in df.loc[user_id, "action"]:
            edge_index.append((user_id, page_id))
    edge_index = torch.tensor(edge_index).t().contiguous()
    user_features = torch.from_numpy(np.vstack(df.features.values)).float()
    page_features = torch.from_numpy(np.vstack(df.page_features.values)).float()
    y = torch.tensor(df.label.values, dtype=torch.int64)
    train_mask = torch.zeros(len(df), dtype=torch.bool)
    valid_mask = torch.zeros(len(df), dtype=torch.bool)
    test_mask = torch.zeros(len(df), dtype=torch.bool)
    train_mask[torch.from_numpy(df.loc[df.flag == "Train"].index.values)] = True
    valid_mask[torch.from_numpy(df.loc[df.flag == "Valid"].index.values)] = True
    test_mask[torch.from_numpy(df.loc[df.flag == "Test"].index.values)] = True
    return {
        "user_features": user_features,
        "page_features": page_features,
        "edge_index": edge_index,
        "labels": y,
        "train_mask": train_mask,
        "valid_mask": valid_mask,
        "test_mask": test_mask,
    }


def get_file_name():
    import inspect
    try:
        # inspect.stack()[0] 是当前函数
        # inspect.stack()[1] 是调用该函数的函数
        caller_frame = inspect.stack()[1]
        return caller_frame.filename.split("\\")[-1].split(".")[0]
    except IndexError:
        # 如果堆栈深度不足，返回 None 或抛出异常
        return None


def show_train_valid_loss_curve(train_losses: list, val_accs: list, test_accs: list, prefix: str) -> None:
    # 绘制训练曲线
    plt.figure(figsize=(12, 4))
    plt.subplot(121)
    plt.plot(train_losses, label='Training Loss')
    plt.title('Training Loss')
    plt.xlabel('Epoch')
    plt.legend()

    plt.subplot(122)
    plt.plot(val_accs, label='Validation Accuracy')
    plt.plot(test_accs, label='Test Accuracy')
    plt.title('Accuracy')
    plt.xlabel('Epoch')
    plt.legend()
    plt.tight_layout()
    plt.savefig(f"./results/pics/{prefix}_training_loss.png", dpi=300)
    plt.show()


# 可视化节点嵌入 (使用测试集)
@torch.no_grad()
def visualize(model: nn.Module, out: torch.Tensor, mask: torch.Tensor, y: torch.Tensor, prefix: str) -> None:
    model.eval()
    z = out[mask].cpu().numpy()
    y = y[mask].cpu().numpy()
    tsne = TSNE(n_components=2, random_state=42)
    z_2d = tsne.fit_transform(z)
    plt.figure(figsize=(8, 8))
    scatter = plt.scatter(z_2d[:, 0], z_2d[:, 1], c=y, cmap='Set1', alpha=0.7)
    plt.legend(*scatter.legend_elements(), title="Classes")
    plt.title('Graph Node Embeddings')
    plt.savefig(f"./results/pics/{prefix}_node_embeddings.png", dpi=300)
    plt.show()


def get_cora_data():
    # 加载数据集
    config = read_yaml("./config.yaml")
    device = config["model_config"]["device"]
    dataset = Planetoid(root='./data/Cora', name='Cora')
    data = dataset[0].to(device)

    # 划分训练/验证/测试集 (使用掩码)
    data.train_mask = torch.zeros(data.num_nodes, dtype=torch.bool)
    data.valid_mask = torch.zeros(data.num_nodes, dtype=torch.bool)
    data.test_mask = torch.zeros(data.num_nodes, dtype=torch.bool)

    # 随机划分：140个训练，500个验证，1000个测试
    indices = torch.randperm(data.num_nodes)
    data.train_mask[indices[:140]] = True
    data.valid_mask[indices[140:640]] = True
    data.test_mask[indices[640:1640]] = True
    return data


def get_diy_data():
    # 设备设置
    config = read_yaml("./config.yaml")
    device = config["model_config"]["device"]
    num_features = config["data_config"]["num_features"]
    num_classes = config["data_config"]["num_classes"]
    model_name = config["model_config"]["graph_model_name"]
    print(f"Using device: {device}")

    # 加载数据集
    dataset = construct_graph_data(df=pd.read_pickle("./dataset.pkl"))
    data = dataset.to(device)
    return data
